/* Copyright 2015 The TensorFlow Authors. All Rights Reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.

https://github.com/tensorflow/tensorflow/blob/master/tensorflow/core/util/sparse/group_iterator.h

Additional modifications made by Huawei Technologies Co., Ltd in 2020-2022.
==============================================================================*/

#ifndef CPU_KERNEL_UTIL_SPARSE_GROUP_ITERATOR_H_
#define CPU_KERNEL_UTIL_SPARSE_GROUP_ITERATOR_H_

#include <vector>
#include "cpu_kernel/utils/eigen_tensor.h"

namespace aicpu {
class Group;  // Predeclare Group for GroupIterable.

// Forward Definition of Sparse Tensor
class GroupIterable {
 public:
  using VarDimArray = std::vector<int64_t>;

  GroupIterable(Tensor *ix, Tensor *vals, int dims, const VarDimArray &group_dims)
      : ix_(ix),
        ix_matrix_(EigenTensor(ix, ix->GetData()).matrix<int64_t>()),
        vals_(vals),
        dims_(dims),
        group_dims_(group_dims.begin(), group_dims.end()) {}

  ~GroupIterable() {}

  class IteratorStep;

  IteratorStep begin() { return IteratorStep(this, 0); }

  IteratorStep at(int64_t loc) { return IteratorStep(this, loc); }

  IteratorStep end() { return IteratorStep(this, ix_->GetTensorShape()->GetDimSize(0)); }

  template <typename TIX>
  inline bool GroupMatches(const TIX &ix, int64_t loc_a, int64_t loc_b) const {
    for (int64_t d : group_dims_) {
      if (ix(loc_a, d) != ix(loc_b, d)) {
        return false;
      }
    }
    return true;
  }

  class IteratorStep {
   public:
    IteratorStep(GroupIterable *iter, int64_t loc) : iter_(iter), location_(loc), next_loc_(location_) {
      UpdateEndOfGroup();
    }

    ~IteratorStep() { iter_ = nullptr; }

    void UpdateEndOfGroup();

    bool operator!=(const IteratorStep &rhs) const;

    bool operator==(const IteratorStep &rhs) const;

    IteratorStep &operator++();

    const IteratorStep operator++(int);

    Group operator*() const;

    int64_t loc() const { return location_; }

   private:
    GroupIterable *iter_;
    int64_t location_;
    int64_t next_loc_;
  };

 private:
  friend class Group;
  Tensor *ix_;
  TTypes<int64_t>::Matrix ix_matrix_;
  Tensor *vals_;
  const int dims_;
  const std::vector<int64_t> group_dims_;
};

// This class is returned when dereferencing a GroupIterable iterator.
// It provides the methods group(), indices(), and values(), which
// provide access into the underlying SparseTensor.
class Group {
 public:
  Group(GroupIterable *iter, int64_t loc, int64_t next_loc) : iter_(iter), loc_(loc), next_loc_(next_loc) {}

  ~Group() { iter_ = NULL; }

  std::vector<int64_t> group() const;

  TTypes<int64_t>::UnalignedConstMatrix indices() const;

  int64_t group_at(size_t index) const {
    const auto &ix_t = iter_->ix_matrix_;
    return ix_t(loc_, index);
  }

  template <typename T>
  typename TTypes<T>::UnalignedVec values() const {
    return typename TTypes<T>::UnalignedVec(&(EigenTensor(iter_->vals_, iter_->vals_->GetData()).vec<T>()(loc_)),
                                            next_loc_ - loc_);
  }

 private:
  GroupIterable *iter_;
  int64_t loc_;
  int64_t next_loc_;
};
}  // namespace aicpu

#endif  // CPU_KERNEL_UTIL_SPARSE_GROUP_ITERATOR_H_
